Skip to content

Add TorchAO disk group offload support#13875

Open
hlky wants to merge 2 commits into
huggingface:mainfrom
hlky:codex/torchao-group-offload
Open

Add TorchAO disk group offload support#13875
hlky wants to merge 2 commits into
huggingface:mainfrom
hlky:codex/torchao-group-offload

Conversation

@hlky

@hlky hlky commented Jun 5, 2026

Copy link
Copy Markdown
Contributor

Split out from #13719 after review feedback to keep the TorchAO safetensors PR focused on save/load.

This follow-up contains only the TorchAO disk group-offload work:

  • allow disk group offload for TorchAO tensor subclasses when torchao safetensors support is available
  • flatten/unflatten TorchAO subclass tensors through safetensors metadata for disk offload
  • add focused TorchAO disk group-offload coverage

Validation so far:

  • python -m ruff check src/diffusers/hooks/group_offloading.py tests/quantization/torchao/test_torchao.py
  • python -m compileall -q src/diffusers/hooks/group_offloading.py tests/quantization/torchao/test_torchao.py

Opening as draft because the parent save/load PR should land first, and this PR may need review after that scope settles.

@hlky hlky marked this pull request as ready for review June 12, 2026 10:46
@hlky

hlky commented Jun 12, 2026

Copy link
Copy Markdown
Contributor Author

@sayakpaul @DN6 This part is ready for review

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

hooks size/L PR with diff > 200 LOC tests

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant